/**
 * 
 */
package edu.umd.clip.lm.model.decoding;

import java.util.*;

import edu.berkeley.nlp.util.Pair;
import edu.umd.clip.lm.util.BinaryPrefix;
import edu.umd.clip.lm.util.ProbMath;
import edu.umd.clip.lm.util.SimpleList;
import edu.umd.clip.lm.factors.Dictionary;
import edu.umd.clip.lm.factors.FactorTuple;
import edu.umd.clip.lm.model.Experiment;
import edu.umd.clip.lm.model.ProbTree;
import edu.umd.clip.lm.model.decoding.VariationalTaggerTrellisSlice.Trie;
import edu.umd.clip.lm.model.decoding.TrellisSlice.SliceItem;
/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class ViterbiTagger extends Tagger {
	
	public static class PathItem {
		final int tag;
		final SliceItem item;
		
		public PathItem(int tag, SliceItem item) {
			this.tag = tag;
			this.item = item;
		}
	}
	
	public static class ViterbiPath implements Comparable<ViterbiPath>, TaggedPath {
		final double logProb;
		final SimpleList<PathItem> tags;
		private static byte tagFactor;
		
		public ViterbiPath(double logProb, SimpleList<PathItem> tags) {
			this.logProb = logProb;
			this.tags = tags;
		}

		public ViterbiPath append(SliceItem item, int tag, double logProb) {
			ViterbiPath viterbiPath = new ViterbiPath(this.logProb + logProb, tags.add(new PathItem(tag, item)));
			return viterbiPath;
		}
		
		@Override
		public int compareTo(ViterbiPath other) {
			return (int) Math.signum(logProb - other.logProb);
		}
		
		public int[] getContext(int ctxSize) {
			int ctx[] = new int[ctxSize];
			SimpleList<PathItem> lst = tags;
			for(int i=0; i<ctxSize; ++i) {
				ctx[i] = lst.head().tag;
				lst = lst.tail();
			}
			return ctx;
		}

		public double getLogProb() {
			return logProb;
		}

		public int[] getOrderedTags() {
			int t[] = new int[tags.length()];
			SimpleList<PathItem> tail = tags;
			for(int i=1; i<=t.length; ++i) {
				// reverse to the proper order
				t[t.length-i] = tail.head().tag;
				tail = tail.tail();
			}
			return t;
		}
		
		public String toString() {
			StringBuilder sb = new StringBuilder();
			Dictionary dict = Experiment.getInstance().getTupleDescription().getDictionary(tagFactor);
			
			int tags[] = getOrderedTags();
			for(int tag : tags) {
				sb.append(dict.getWord(tag));
				sb.append(' ');
			}
			sb.append('(');
			sb.append(logProb);
			sb.append(')');
			return sb.toString();
		}
	}
	
	public ViterbiTagger(Decoder decoder, int beamSize, String tagFactor) {
		super(decoder, beamSize, tagFactor);
		ViterbiPath.tagFactor = tagFactorIdx;
	}
	
	public Pair<SortedSet<? extends TaggedPath>, Double> evaluateSentence(Collection<FactorTuple> sentence) {
		SentenceContext ctx = decoder.createContext();
		ViterbiTaggerContext taggerContext = new ViterbiTaggerContext(tags);
		final int ctxSize = decoder.getForest().getOrder() - 1;
		//HashMap<TaggerContext.Context,List<ViterbiPath>> currentPathsHash = taggerContext.currentPathsHash;
		//TreeSet<ViterbiPath> currentPaths = taggerContext.currentPaths;
		
		ViterbiPath initPath = null;
		
		{
			SimpleList<PathItem> seq = new SimpleList<PathItem>();
			
			for(int i=0; i<ctxSize; ++i) {
				SliceItem item = ctx.getSliceByOffset(i+1-ctxSize).getElements().iterator().next();
				seq = seq.add(new PathItem(Dictionary.getStart(),item));
			}
			initPath = new ViterbiPath(0.0, seq);
		}
		taggerContext.currentPaths.add(initPath);
		
		{
			List<ViterbiPath> lst = new LinkedList<ViterbiPath>();
			lst.add(initPath);
			//taggerContext.currentPathsHash.put(new ViterbiTaggerContext.Context(initPath.getContext(ctxSize)), lst);
		}
		
		long overtMask = Experiment.getInstance().getTupleDescription().getOvertFactorsMask();
		for(FactorTuple tuple : sentence) {
			tuple = new FactorTuple(tuple.getBits() & overtMask);
			ctx = ctx.addWord(tuple.getBits());
			decoder.partition(ctx);
			if (decoder.getDebug() > 0) {
				decoder.debugContext(ctx);
				if (decoder.getDebug() > 1) {
					ctx.getTrellis().dumpTrellis(String.format("trellis-tagger-%d-%s.dot", 
							ctx.getTrellis().getCurrentSlice().getTime(), tuple.toStringNoNull().replace('/', ' ')));
				}

			}
			if (ctx.getCurrentSlice().getTime() <= ctxSize * 2) {
				/*
				SimpleList<PathItem> seq;
				SliceItem item = ctx.getCurrentSlice().getPrevious().getElements().iterator().next();
				if (initPath == null) {
					seq = new SimpleList<PathItem>();
					seq = seq.add(new PathItem(Dictionary.getStart(), item));
					initPath = new ViterbiPath(0.0, seq);
				} else {
					initPath = initPath.append(item, Dictionary.getStart(), 0.0);
				}
				*/
				continue;
			}
			TrellisSlice slice = ctx.getSliceByOffset(-ctxSize); 
			/*
			if (ctx.getCurrentSlice().getTime() == ctxSize+2) {
				taggerContext.currentPaths.add(initPath);
			}
			*/
			long start = System.nanoTime();
			
			splitElements(slice);
			
			long end = System.nanoTime();
			long updatePathsDelta = end - start;
			
			updatePathsTime += updatePathsDelta;
			
			start = end;
			
			updatePaths(taggerContext, slice);
			
			end = System.nanoTime();
			long makingTaggerSliceDelta = end - start;
			
			makingTaggerSliceTime += makingTaggerSliceDelta;
			
			if (decoder.getDebug() > 0) {
				System.out.printf("tagger [%d]: %d paths, update: %fs, slice: %fs\n",
						slice.getTime(),
						taggerContext.currentPaths.size(),
						(double)updatePathsDelta/1000000000,
						(double)makingTaggerSliceDelta/1000000000);
				//printSlice(slice);
			}
		}
		// do the sentence terminator
		for(int offset=-ctxSize+1; offset<=0; ++offset) {
			long start = System.nanoTime();
			TrellisSlice slice = ctx.getSliceByOffset(offset); 
			splitElements(slice);
			
			long end = System.nanoTime();
			updatePathsTime += end - start;
			start = end;
			
			
			updatePaths(taggerContext, slice);
			
			end = System.nanoTime();
			makingTaggerSliceTime += end - start;
	
			if (decoder.getDebug() > 0) {
				//printSlice(slice);
			}
		}
		return new Pair<SortedSet<? extends TaggedPath>, Double>(taggerContext.currentPaths, ctx.getLogProb());
	}
	
	private void updatePaths(ViterbiTaggerContext taggerContext, TrellisSlice slice) {
		final int ctxSize = decoder.getForest().getOrder() - 1;

		TreeSet<ViterbiPath> viterbiPaths = new TreeSet<ViterbiPath>();
		//HashMap<TaggerContext.Context, List<ViterbiPath>> pathsHash = new HashMap<TaggerContext.Context, List<ViterbiPath>>(taggerContext.currentPathsHash.size());
		
		double minProb = Double.NEGATIVE_INFINITY;
		//for(Map.Entry<TaggerContext.Context, List<ViterbiPath>> entry : taggerContext.currentPathsHash.entrySet()) {
		for(ViterbiPath prevPath : taggerContext.currentPaths.descendingSet()) {
			if (viterbiPaths.size() >= beamSize && prevPath.logProb < minProb) {
				// pruned
				break;
			}

			PathItem head = prevPath.tags.head();
			SliceItem item = head.item;
			
			if (item == null) {
				// the first element
				SliceItem nextItem = slice.getElements().iterator().next();
				double logProb = Math.log10(nextItem.tree.getTotalProb());
				int tag = taggerContext.getTag(nextItem.tree.getPrefix());
				ViterbiPath newPath = prevPath.append(nextItem, tag, logProb);
				viterbiPaths.add(newPath);
				break;
			}
			
			List<SliceItem> fwdLinks = slice.prev2curr.get(item);
			if (fwdLinks != null) {
				TrellisSlice afterNextSlice = slice.getNext();
				if (afterNextSlice != null) afterNextSlice = afterNextSlice.getNext();
				
				for(SliceItem nextItem : fwdLinks) {
					double logProb = Math.log10(nextItem.tree.getTotalProb());
					if (viterbiPaths.size() >= beamSize && logProb + prevPath.logProb < minProb) {
						// pruned
						continue;
					}
					if (ctxSize > 1 && afterNextSlice != null && afterNextSlice.prev2curr.containsKey(nextItem)) {
						// a dead end
						continue;
					}
					
					int tag = taggerContext.getTag(nextItem.tree.getPrefix());
					ViterbiPath newPath = prevPath.append(nextItem, tag, logProb);
					viterbiPaths.add(newPath);
					if (viterbiPaths.size() > beamSize) {
						viterbiPaths.pollFirst();
						minProb = viterbiPaths.first().logProb;
					}
				}
			}
		}
		taggerContext.currentPaths = viterbiPaths;
		//taggerContext.currentPathsHash = pathsHash;
	}
	
	private void splitElements(TrellisSlice slice) {
		SliceItem elements[] = slice.getElements().toArray(new SliceItem[slice.getElements().size()]);
		
		for(SliceItem sliceItem : elements) {
			BinaryPrefix prefix = sliceItem.tree.getPrefix();
			for(int tag = 0; tag < tags.length; ++tag) {
				BinaryPrefix tagPrefix = tags[tag];
				if (tagPrefix == null) continue;
				
				if (tagPrefix.isPrefixOf(prefix)) {
					break;
				}
				if (prefix.isPrefixOf(tagPrefix)) {
					//ProbTree tmpBackup = (ProbTree) sliceItem.tree.clone();
					ProbTree probTree = sliceItem.tree.cut(tagPrefix);
					if (probTree == sliceItem.tree) {
						sliceItem.tree = probTree;
						break; // no need to check other tags
					}
					if (probTree != null) {
						SliceItem newItem = new SliceItem(probTree, sliceItem.getPrevProb(), (ClusterId)sliceItem.clusterid.clone());
						slice.split(sliceItem, newItem);
						//System.out.printf("tagger: cutting %s off %s\n", newItem, sliceItem);
					}
				}
			}
		}		
	}
	
	private void printSlice(VariationalTaggerTrellisSlice slice) {
		int ctxSize = decoder.getForest().getOrder() - 1;
		int ctx[] = new int[ctxSize];
		if (slice.trie.subTries != null) {
			for(Trie subTrie: slice.trie.subTries) {
				ctx[0] = subTrie.tag;
				printSliceIter(subTrie, ctxSize, ctx);
			}
		}
	}

	private void printSliceIter(Trie trie, int stepsLeft, int ctx[]) {
		if (stepsLeft > 1) {
			if (trie.subTries != null) {
				for(Trie subTrie: trie.subTries) {
					ctx[stepsLeft-1] = subTrie.tag;
					printSliceIter(subTrie, stepsLeft-1, ctx);
				}
			}
		} else {
			Dictionary dict = Experiment.getInstance().getTupleDescription().getDictionary(ViterbiPath.tagFactor);

			for(int i=ctx.length-1; i>=0; --i) {
				System.out.print(dict.getWord(ctx[i]));
				System.out.print(' ');
			}
			System.out.printf("%f", Math.log10(trie.prob));
			System.out.print(" -> ( ");
			if (trie.subTries != null) {
				for(Trie subTrie : trie.subTries) {
					System.out.printf("%s:%f ", dict.getWord(subTrie.tag), Math.log10(subTrie.prob));
				}
			}
			System.out.print(")\n");
		}
	}
}
